import asyncio
import websockets
import json
import os
import base64
from datetime import datetime


class ChatServer:
    def __init__(self):
        self.clients = {}  # 存储连接的客户端 {websocket: {"user_id": ...}}
        self.rooms = {}    # 存储房间信息 {"room_id": [client1, client2, ...]}

    async def register_client(self, websocket, user_data):
        """注册新客户端"""
        user_id = user_data.get("user_id")
        room_id = user_data.get("room_id", "default")

        # 将客户端添加到全局客户端列表
        self.clients[websocket] = {
            "user_id": user_id,
            "room_id": room_id
        }

        # 创建房间（如果不存在）
        if room_id not in self.rooms:
            self.rooms[room_id] = []

        # 将客户端添加到房间
        self.rooms[room_id].append(websocket)

        # 通知房间内其他用户有新用户加入
        join_message = {
            "type": "system",
            "message": f"用户 {user_id} 加入了房间",
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "sender": "System"
        }
        await self.broadcast_to_room(room_id, join_message, websocket)

        # 发送欢迎消息给新用户
        welcome_message = {
            "type": "system",
            "message": f"欢迎用户 {user_id} 加入房间 {room_id}",
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "sender": "System"
        }
        await websocket.send(json.dumps(welcome_message))

        # 发送当前房间用户列表
        await self.send_user_list(room_id)

    async def unregister_client(self, websocket):
        """注销客户端"""
        if websocket in self.clients:
            client_info = self.clients[websocket]
            user_id = client_info["user_id"]
            room_id = client_info["room_id"]

            # 从客户端列表中移除
            del self.clients[websocket]

            # 从房间中移除
            if room_id in self.rooms and websocket in self.rooms[room_id]:
                self.rooms[room_id].remove(websocket)

            # 通知房间内其他用户该用户已离开
            leave_message = {
                "type": "system",
                "message": f"用户 {user_id} 离开了房间",
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                "sender": "System"
            }
            await self.broadcast_to_room(room_id, leave_message, None)

            # 更新房间用户列表
            await self.send_user_list(room_id)

    async def broadcast_to_room(self, room_id, message, exclude_websocket=None):
        """向房间内所有客户端广播消息"""
        if room_id in self.rooms:
            # 并发发送消息给所有客户端以提高性能
            tasks = []
            for client in self.rooms[room_id][:]:  # 使用切片复制避免修改列表时出错
                if client != exclude_websocket:
                    try:
                        # 创建发送任务而不是直接等待
                        task = asyncio.create_task(client.send(json.dumps(message)))
                        tasks.append(task)
                    except websockets.exceptions.ConnectionClosed:
                        # 如果连接已关闭，则移除客户端
                        await self.unregister_client(client)

            # 并发执行所有发送任务
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)

    async def send_user_list(self, room_id):
        """发送房间用户列表"""
        if room_id in self.rooms:
            user_ids = []  # 改为存储用户ID而不是用户名
            for client in self.rooms[room_id]:
                if client in self.clients:
                    user_ids.append(self.clients[client]["user_id"])  # 使用用户ID

            user_list_message = {
                "type": "user_list",
                "users": user_ids,  # 发送用户ID列表
                "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            }
            await self.broadcast_to_room(room_id, user_list_message)

    async def handle_message(self, websocket, message_data):
        """处理接收到的消息"""
        message_type = message_data.get("type")
        # print(f"Received message type: {message_type}: {message_data}")

        if message_type == "register":
            # 用户注册
            await self.register_client(websocket, message_data)

        elif message_type == "text":
            # 文本消息
            if websocket in self.clients:
                client_info = self.clients[websocket]
                text_message = {
                    "type": "text",
                    "message": message_data.get("message"),
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "sender": client_info["user_id"],  # 使用用户ID而不是用户名
                    "sender_id": client_info["user_id"]
                }
                await self.broadcast_to_room(client_info["room_id"], text_message, websocket)

        elif message_type == "file":
            # 文件传输
            if websocket in self.clients:
                client_info = self.clients[websocket]
                file_message = {
                    "type": "file",
                    "filename": message_data.get("filename"),
                    "filedata": message_data.get("filedata"),  # base64编码的文件数据
                    "filesize": message_data.get("filesize"),
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                    "sender": client_info["user_id"],  # 使用用户ID而不是用户名
                    "sender_id": client_info["user_id"]
                }
                await self.broadcast_to_room(client_info["room_id"], file_message, websocket)

    async def handle_client(self, websocket, path):
        """处理客户端连接"""
        try:
            async for message in websocket:
                try:
                    message_data = json.loads(message)
                    await self.handle_message(websocket, message_data)
                except json.JSONDecodeError:
                    error_message = {
                        "type": "error",
                        "message": "无效的消息格式",
                        "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                    }
                    await websocket.send(json.dumps(error_message))
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            await self.unregister_client(websocket)


# 创建服务器实例
chat_server = ChatServer()

# 启动服务器
start_server = websockets.serve(chat_server.handle_client, "0.0.0.0", 18229)

print("聊天服务器启动在 ws://0.0.0.0:18229")

# 运行事件循环
asyncio.get_event_loop().run_until_complete(start_server)
asyncio.get_event_loop().run_forever()